/*================================================================
*   文件名称：main.cpp
*   描    述：
*
* ================================================================*/


#include <iostream>
#include <vector>

using namespace std;

class Solution {
public:
    /**
     * @param matrix: a matrix
     * @return: return how many square submatrices have all ones
     */
    int countSquares(vector<vector<int>> &matrix) {
        // write your code here
        int m = matrix.size();
        if(m == 0) {
            return 0;
        }

        int n = matrix[0].size();
        if(n == 0) {
            return 0;
        }

        int ans = 0;
        vector<vector<int>> dp(m, vector<int>(n, 0));
        for(int i = 0; i < m; i++) {
            for(int j = 0; j < n; j++) {
                if(i == 0 || j == 0) {
                    dp[i][j] = matrix[i][j];
                } else if(matrix[i][j] == 1) {
                    int minR = min(dp[i-1][j], dp[i][j-1]);
                    dp[i][j] = min(dp[i-1][j-1], minR) + 1;
                }

                ans += dp[i][j];
            }
        }

        return ans;
    }
};
